package com.lagou.filter;

import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.LinkedList;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
@Order(0)
public class IPFilter implements GatewayFilter {

    /**
     * 存储每个ip请求的访问量 key = ip; value = 访问时间
     */
    private volatile static Map<String, LinkedList<Long>> ipMap = new ConcurrentHashMap<>();

    /**
     * 最大请求数
     */
    @Value("${prevent.X}")
    private Integer x;

    /**
     * 多少分钟内
     */
    @Value("${prevent.Y}")
    private Integer y;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        ServerHttpResponse response = exchange.getResponse();

        String ip = request.getRemoteAddress().getAddress().getHostAddress();
        ip = ip.equals("0:0:0:0:0:0:0:1") ? "127.0.0.1" : ip;
        LinkedList<Long> times = ipMap.get(ip);
        if (times == null) {
            times = new LinkedList<>();
            times.add(System.currentTimeMillis());
            ipMap.put(ip, times);
        } else {
            times.add(System.currentTimeMillis());
        }

        //根据请求 y 限制，获取请求总数
        int reqNum = getReqNum(ipMap.get(ip));
        if (reqNum > x) {
            //超过请求阈值
            response.setStatusCode(HttpStatus.SEE_OTHER);
            DataBuffer wrap = response.bufferFactory().wrap("reject!".getBytes());
            return response.writeWith(Mono.justOrEmpty(wrap));
        }

        //合法请求放行
        return chain.filter(exchange);
    }

    private int getReqNum(LinkedList<Long> times) {
        int l = y * 60 * 1000;
        int reqNum =0;
        for (Long time : times) {
            if (System.currentTimeMillis() - l <= time ) {
                //未过期 +1
                reqNum++;
            }
        }
        return reqNum;
    }
}
